"""
总共10类
每类选择500张图像
每张图预计做10种数据增强变换

总计 10*500*10 = 50000
"""
import torch
import torchvision
import random
import cv2
import numpy as np
import torch.nn as nn
import torchvision.transforms as transforms
import torch.nn.functional as F
from test import load_model
from numpy import linalg as LA
import math
import time
import sys
from PIL import Image
import torch.utils.data as data


class OriDataset(torch.utils.data.Dataset):
    def __init__(self, transform):
        images = np.load('../data.npy')
        labels = np.load('../label.npy')
        assert labels.min() >= 0
        assert images.dtype == np.uint8
        self.images = [Image.fromarray(x) for x in images]
        self.labels = labels / labels.sum(axis=1, keepdims=True)  # normalize
        self.labels = self.labels.astype(np.float32)
        self.transform = transform

    def __getitem__(self, index):
        image, label = self.images[index], self.labels[index]
        image = self.transform(image)
        return image, label

    def __len__(self):
        return len(self.labels)


if __name__ == '__main__':
    print('hello world!')
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    models = []

    resnet50 = load_model('F:\\AliTianChi\\ali_8_attack\\resnet50.pth.tar', ResTag=True)
    resnet50.eval()
    models.append(resnet50)

    densenet121 = load_model('F:\\AliTianChi\\ali_8_attack\\densenet121.pth.tar', ResTag=False)
    densenet121.eval()
    models.append(densenet121)

    transform_train = transforms.Compose([
        transforms.ToTensor(),
    ])
    normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))

    trainset = OriDataset(transform=transform_train)
    trainloader = data.DataLoader(trainset, batch_size=1, shuffle=False, num_workers=0)
    count_list = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    data_select = []
    label_select = []
    for (inputs, soft_labels) in trainloader:
        inputs = inputs.to(device)
        targets = soft_labels.argmax(dim=-1)
        targets = targets.to(device)
        if count_list[targets.item()] >= 2500:
            continue

        r1 = resnet50(normalize(inputs))
        _, r1_index = torch.topk(r1, 1)
        if r1_index.item() != targets.item():
            continue

        r2 = densenet121(normalize(inputs))
        _, r2_index = torch.topk(r2, 1)
        if r2_index.item() != targets.item():
            continue

        # 都通过了，记录一下
        count_list[targets.item()] += 1
        inputs = inputs[0]
        inputs = inputs * 255
        inputs = inputs.cpu().detach().numpy()
        inputs = inputs.transpose(1, 2, 0)
        data_select.append(inputs)
        soft_labels = soft_labels.numpy()
        label_select.append(soft_labels[0])
        # if count_list[targets.item()] == 500

    data_select = np.asarray(data_select, dtype=np.uint8)
    label_select = np.asarray(label_select)

    np.save('data_select.npy', data_select)
    np.save('label_select.npy', label_select)
    exit()